{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "from __future__ import print_function\n",
    "import sys\n",
    "\n",
    "import argparse\n",
    "import os\n",
    "import shutil\n",
    "import time\n",
    "import random\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.parallel\n",
    "import torch.backends.cudnn as cudnn\n",
    "import torch.optim as optim\n",
    "import torch.utils.data as data\n",
    "import torchvision.transforms as transforms\n",
    "import torchvision.datasets as datasets\n",
    "import torchvision.models as models\n",
    "import models.imagenet as customized_models\n",
    "from flops_counter import get_model_complexity_info\n",
    "from PIL import ImageFile\n",
    "from PIL import Image\n",
    "ImageFile.LOAD_TRUNCATED_IMAGES = True\n",
    "\n",
    "from utils import Bar, Logger, AverageMeter, accuracy, mkdir_p\n",
    "\n",
    "# for servers to immediately record the logs\n",
    "def flush_print(func):\n",
    "    def new_print(*args, **kwargs):\n",
    "        func(*args, **kwargs)\n",
    "        sys.stdout.flush()\n",
    "    return new_print\n",
    "print = flush_print(print)\n",
    "\n",
    "\n",
    "# Models\n",
    "default_model_names = sorted(name for name in models.__dict__\n",
    "    if name.islower() and not name.startswith(\"__\")\n",
    "    and callable(models.__dict__[name]))\n",
    "\n",
    "customized_models_names = sorted(name for name in customized_models.__dict__\n",
    "    if name.islower() and not name.startswith(\"__\")\n",
    "    and callable(customized_models.__dict__[name]))\n",
    "\n",
    "for name in customized_models.__dict__:\n",
    "    if name.islower() and not name.startswith(\"__\") and callable(customized_models.__dict__[name]):\n",
    "        models.__dict__[name] = customized_models.__dict__[name]\n",
    "\n",
    "model_names = default_model_names + customized_models_names\n",
    "\n",
    "# Parse arguments\n",
    "parser = argparse.ArgumentParser(description='PyTorch ImageNet Training')\n",
    "\n",
    "# Datasets\n",
    "parser.add_argument('-d', '--data', default='./data', type=str)           #default is \"path to dataset\"\n",
    "parser.add_argument('-j', '--workers', default=4, type=int, metavar='N',\n",
    "                    help='number of data loading workers (default: 4)')\n",
    "# Optimization options\n",
    "parser.add_argument('--epochs', default=10, type=int, metavar='N',      #default is 90\n",
    "                    help='number of total epochs to run')\n",
    "parser.add_argument('--start-epoch', default=0, type=int, metavar='N',\n",
    "                    help='manual epoch number (useful on restarts)')\n",
    "parser.add_argument('--train-batch', default=20, type=int, metavar='N', #default is 256\n",
    "                    help='train batchsize (default: 256)')\n",
    "parser.add_argument('--test-batch', default=20, type=int, metavar='N',  #default is 200\n",
    "                    help='test batchsize (default: 200)')\n",
    "parser.add_argument('--lr', '--learning-rate', default=0.0005, type=float,\n",
    "                    metavar='LR', help='initial learning rate')\n",
    "parser.add_argument('--drop', '--dropout', default=0, type=float,\n",
    "                    metavar='Dropout', help='Dropout ratio')\n",
    "parser.add_argument('--schedule', type=int, nargs='+', default=[150, 225],\n",
    "                        help='Decrease learning rate at these epochs.')\n",
    "parser.add_argument('--gamma', type=float, default=0.1, help='LR is multiplied by gamma on schedule.')\n",
    "parser.add_argument('--momentum', default=0.9, type=float, metavar='M',\n",
    "                    help='momentum')\n",
    "parser.add_argument('--weight-decay', '--wd', default=1e-4, type=float,\n",
    "                    metavar='W', help='weight decay (default: 1e-4)')\n",
    "# Checkpoints\n",
    "parser.add_argument('-c', '--checkpoint', default='./checkpoint/old_resnet50', type=str, metavar='PATH',   #deault is checkpoint\n",
    "                    help='path to save checkpoint (default: checkpoint)')\n",
    "parser.add_argument('--resume', default='../pretrain/old_resnet50.pth.tar', type=str, metavar='PATH',  #default is none    '../pretrain/old_resnet50.pth.tar'\n",
    "                    help='path to latest checkpoint (default: none)')\n",
    "# Architecture\n",
    "parser.add_argument('--modelsize', '-ms', metavar='large', default='large', \\\n",
    "                    choices=['large', 'small'], \\\n",
    "                    help = 'model_size affects the data augmentation, please choose:' + \\\n",
    "                           ' large or small ')\n",
    "parser.add_argument('--arch', '-a', metavar='ARCH', default='old_resnet50',\n",
    "                    choices=model_names,\n",
    "                    help='model architecture: ' +\n",
    "                        ' | '.join(model_names) +\n",
    "                        ' (default: resnet18)')\n",
    "parser.add_argument('--depth', type=int, default=29, help='Model depth.')\n",
    "parser.add_argument('--cardinality', type=int, default=32, help='ResNet cardinality (group).')\n",
    "parser.add_argument('--base-width', type=int, default=4, help='ResNet base width.')\n",
    "parser.add_argument('--widen-factor', type=int, default=4, help='Widen factor. 4 -> 64, 8 -> 128, ...')\n",
    "# Miscs\n",
    "parser.add_argument('--manualSeed', type=int, help='manual seed')\n",
    "parser.add_argument('-e', '--evaluate', dest='evaluate', action='store_true',\n",
    "                    help='evaluate model on validation set')\n",
    "parser.add_argument('--pretrained', dest='pretrained', action='store_true',\n",
    "                    help='use pre-trained model')\n",
    "#Device options\n",
    "parser.add_argument('--gpu-id', default='0', type=str,\n",
    "                    help='id(s) for CUDA_VISIBLE_DEVICES')\n",
    "parser.add_argument('-o','--output',dest='output',action='store_true',help='output_classification')\n",
    "# parser.add_argument('-t','--test',dest='test_mode',action='test_img',help='pic_demo')\n",
    "\n",
    "args = parser.parse_args()\n",
    "state = {k: v for k, v in args._get_kwargs()}\n",
    "\n",
    "# Use CUDA\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_id\n",
    "use_cuda = torch.cuda.is_available()\n",
    "\n",
    "# Random seed\n",
    "if args.manualSeed is None:\n",
    "    args.manualSeed = random.randint(1, 10000)\n",
    "random.seed(args.manualSeed)\n",
    "torch.manual_seed(args.manualSeed)\n",
    "if use_cuda:\n",
    "    torch.cuda.manual_seed_all(args.manualSeed)\n",
    "\n",
    "best_acc = 0  # best test accuracy\n",
    "\n",
    "def main():\n",
    "    global best_acc\n",
    "    start_epoch = args.start_epoch  # start from epoch 0 or last checkpoint epoch\n",
    "\n",
    "    if not os.path.isdir(args.checkpoint):\n",
    "        mkdir_p(args.checkpoint)\n",
    "\n",
    "    # Data loading code\n",
    "    traindir = os.path.join(args.data, 'train')\n",
    "    valdir = os.path.join(args.data, 'val')\n",
    "    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],\n",
    "                                     std=[0.229, 0.224, 0.225])\n",
    "\n",
    "    data_aug_scale = (0.08, 1.0) if args.modelsize == 'large' else (0.2, 1.0)\n",
    "\n",
    "    train_loader = torch.utils.data.DataLoader(\n",
    "        datasets.ImageFolder(traindir, transforms.Compose([\n",
    "            transforms.RandomResizedCrop(224, scale = data_aug_scale),\n",
    "            transforms.RandomHorizontalFlip(),\n",
    "            # transforms.RandomChoice([transforms.ColorJitter(brightness=random.uniform(0.8,1.2),contrast=random.uniform(0.8,1.2),\n",
    "            #                                                 saturation=random.uniform(0.8,1.2)),\n",
    "            #                         transforms.RandomGrayscale(p=0.1),\n",
    "            #                          ]),\n",
    "            transforms.ToTensor(),\n",
    "            normalize,\n",
    "        ])),\n",
    "        batch_size=args.train_batch, shuffle=True,\n",
    "        num_workers=args.workers, pin_memory=True)\n",
    "\n",
    "    val_loader = torch.utils.data.DataLoader(\n",
    "        datasets.ImageFolder(valdir, transforms.Compose([\n",
    "            transforms.Scale(256),\n",
    "            transforms.CenterCrop(224),\n",
    "            transforms.ToTensor(),\n",
    "            normalize,\n",
    "        ])),\n",
    "        batch_size=args.test_batch, shuffle=True,\n",
    "        num_workers=args.workers, pin_memory=True)\n",
    "\n",
    "    # create model\n",
    "    if args.pretrained:\n",
    "        print(\"=> using pre-trained model '{}'\".format(args.arch))\n",
    "        model = models.__dict__[args.arch](pretrained=True)\n",
    "    elif 'resnext' in args.arch:\n",
    "        model = models.__dict__[args.arch](\n",
    "                    baseWidth=args.base_width,\n",
    "                    cardinality=args.cardinality,\n",
    "                )\n",
    "    else:\n",
    "        print(\"=> creating model '{}'\".format(args.arch))\n",
    "        model = models.__dict__[args.arch]()\n",
    "\n",
    "    flops, params = get_model_complexity_info(model, (224, 224), as_strings=False, print_per_layer_stat=False)\n",
    "    print('Flops:  %.3f' % (flops / 1e9))\n",
    "    print('Params: %.2fM' % (params / 1e6))\n",
    "\n",
    "\n",
    "    if args.arch.startswith('alexnet') or args.arch.startswith('vgg'):\n",
    "        model.features = torch.nn.DataParallel(model.features)\n",
    "        model.cuda()\n",
    "    else:\n",
    "        model = torch.nn.DataParallel(model).cuda()\n",
    "\n",
    "    cudnn.benchmark = True\n",
    "\n",
    "    # define loss function (criterion) and optimizer\n",
    "    criterion = nn.CrossEntropyLoss().cuda()\n",
    "    optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)\n",
    "    # Resume\n",
    "    title = 'ImageNet-' + args.arch\n",
    "    # cmp = model.module.fc.weight\n",
    "\n",
    "    if args.resume:\n",
    "        # Load checkpoint.\n",
    "        print('==> Resuming from checkpoint..', args.resume)\n",
    "        assert os.path.isfile(args.resume), 'Error: no checkpoint directory found!'\n",
    "        args.checkpoint = os.path.dirname(args.resume)\n",
    "        checkpoint = torch.load(args.resume)\n",
    "\n",
    "        best_acc = checkpoint['best_acc']\n",
    "        start_epoch = checkpoint['epoch']\n",
    "        # model may have more keys\n",
    "        t = model.state_dict()\n",
    "        c = checkpoint['state_dict']\n",
    "\n",
    "        if \"checkpoint\" not in  args.resume:\n",
    "            tmp1 = c['module.fc.weight'][0:2]\n",
    "            tmp1[1] = c['module.fc.weight'][627]\n",
    "            c['module.fc.weight'] = tmp1\n",
    "            tmp2 = c['module.fc.bias'][0:2]\n",
    "            c['module.fc.bias'] = tmp2\n",
    "            #c['module.fc.weight']*=0\n",
    "        model.module.fc = nn.Linear(2048,2,True)\n",
    "        model.cuda()\n",
    "        flag = True\n",
    "        for k in t:\n",
    "            if k not in c:\n",
    "                print('not in loading dict! fill it', k, t[k])\n",
    "                c[k] = t[k]\n",
    "                flag = False\n",
    "        model.load_state_dict(c)\n",
    "        if flag:\n",
    "            print('optimizer load old state')\n",
    "            optimizer.load_state_dict(checkpoint['optimizer'])\n",
    "        else:\n",
    "            print('new optimizer !')\n",
    "        logger = Logger(os.path.join(args.checkpoint, 'log.txt'), title=title, resume=True)\n",
    "    else:\n",
    "        logger = Logger(os.path.join(args.checkpoint, 'log.txt'), title=title)\n",
    "        logger.set_names(['Learning Rate', 'Train Loss', 'Valid Loss', 'Train Acc.', 'Valid Acc.'])\n",
    "\n",
    "\n",
    "    if args.evaluate:\n",
    "        print('\\nEvaluation only')\n",
    "        test_loss, test_acc = test(val_loader, model, criterion, start_epoch, use_cuda)\n",
    "        print(' Test Loss:  %.8f, Test Acc:  %.2f' % (test_loss, test_acc))\n",
    "        return 24000\n",
    "\n",
    "\n",
    "\n",
    "    if args.output:\n",
    "        normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],\n",
    "                                         std=[0.229, 0.224, 0.225])\n",
    "        val_transforms = transforms.Compose([transforms.Scale(256),\n",
    "                                             transforms.CenterCrop(224),\n",
    "                                             transforms.ToTensor(),\n",
    "                                             normalize])\n",
    "        demo_path = os.path.join(args.data,\"demo\")\n",
    "        img_dataset = datasets.ImageFolder(demo_path,transform=val_transforms)\n",
    "        TP, TN, FP, FN , correct = 0, 0, 0, 0, 0\n",
    "        for idx , (img_path, img_label) in enumerate(img_dataset.imgs):\n",
    "            img_name = img_path.split('/')[-1]\n",
    "            img = Image.open(img_path)\n",
    "            img_out = val_transforms(img)\n",
    "            img_out = torch.unsqueeze(img_out, 0)\n",
    "            with torch.no_grad():\n",
    "                img_out = img_out.to('cuda')\n",
    "                feat = model(img_out)\n",
    "                feat = torch.squeeze(feat,0)\n",
    "                if feat[0] >= feat[1]:\n",
    "                    save_path = \"./data/demo_result/car/\"\n",
    "                    if img_label == 0:\n",
    "                        correct += 1\n",
    "                        TN += 1\n",
    "                        img.save(save_path+img_name)\n",
    "                    else:\n",
    "                        FN += 1\n",
    "                        print(FN,\"th motorcycle is wrongly considered as car.\")\n",
    "                        img.save(save_path+str(FN)+\".jpg\")\n",
    "                else:\n",
    "                    save_path = \"./data/demo_result/motorcycle/\"\n",
    "                    if img_label == 0:\n",
    "                        FP += 1\n",
    "                        print(FP, \"th car is wrongly considered as motorcycle.\")\n",
    "                        img.save(save_path+str(FP)+\".jpg\")\n",
    "                    else:\n",
    "                        correct+=1\n",
    "                        TP += 1\n",
    "                        img.save(save_path+img_name)\n",
    "\n",
    "        print(\"The number of correctly classified pic is  \",correct)\n",
    "        print(\"The acc is {:.4f}\".format(correct/len(img_dataset)))\n",
    "        print(\"The precision is {:.4f}\".format(TP/(TP+FP)))\n",
    "        print(\"The recall is {:.4f}\".format(TP/(TP+FN)))\n",
    "        return\n",
    "\n",
    "\n",
    "\n",
    "    # ignored_params = list(map(id,model.module.fc.parameters()))\n",
    "    # base_params = filter(lambda p: id(p) not in ignored_params, model.parameters())\n",
    "    # params_list = [{'params':base_params,'lr':args.lr},]\n",
    "    # params_list.append({'params':model.module.fc.parameters(),'lr':0})\n",
    "    #\n",
    "    #\n",
    "    # optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)\n",
    "    # Train and val\n",
    "    for epoch in range(0, args.epochs):               #defult is for epoch in range(start_epoch, args.epochs):\n",
    "        print(model.module.fc.weight[0][0])\n",
    "        print(model.module.fc.weight[0][1000])\n",
    "        print(model.module.fc.weight[1][2000])\n",
    "\n",
    "        adjust_learning_rate(optimizer, epoch)\n",
    "\n",
    "        print('\\nEpoch: [%d | %d] LR: %f' % (epoch + 1, args.epochs, state['lr']))\n",
    "\n",
    "        train_loss, train_acc = train(train_loader, model, criterion, optimizer, epoch, use_cuda)\n",
    "        test_loss, test_acc = test(val_loader, model, criterion, epoch, use_cuda)\n",
    "\n",
    "        # append logger file\n",
    "        logger.append([state['lr'], train_loss, test_loss, train_acc, test_acc])\n",
    "\n",
    "        # save model\n",
    "        is_best = test_acc > best_acc\n",
    "        best_acc = max(test_acc, best_acc)\n",
    "        save_checkpoint({\n",
    "                'epoch': epoch + 1,\n",
    "                'state_dict': model.state_dict(),\n",
    "                'acc': test_acc,\n",
    "                'best_acc': best_acc,\n",
    "                'optimizer' : optimizer.state_dict(),\n",
    "            }, is_best, checkpoint=args.checkpoint)\n",
    "\n",
    "    logger.close()\n",
    "\n",
    "    print('Best acc:')\n",
    "    print(best_acc)\n",
    "\n",
    "def train(train_loader, model, criterion, optimizer, epoch, use_cuda):\n",
    "    # switch to train mode\n",
    "    model.train()\n",
    "    torch.set_grad_enabled(True)\n",
    "\n",
    "    batch_time = AverageMeter()\n",
    "    data_time = AverageMeter()\n",
    "    losses = AverageMeter()\n",
    "    top1 = AverageMeter()\n",
    "    top5 = AverageMeter()\n",
    "    end = time.time()\n",
    "\n",
    "    bar = Bar('Processing', max=len(train_loader))\n",
    "    show_step = len(train_loader) // 10\n",
    "    for batch_idx, (inputs, targets) in enumerate(train_loader):\n",
    "        batch_size = inputs.size(0)\n",
    "        if batch_size < args.train_batch:\n",
    "            continue\n",
    "        # measure data loading time\n",
    "        data_time.update(time.time() - end)\n",
    "\n",
    "        if use_cuda:\n",
    "            inputs, targets = inputs.cuda(), targets.cuda()\n",
    "        inputs, targets = torch.autograd.Variable(inputs), torch.autograd.Variable(targets)\n",
    "\n",
    "        # compute output\n",
    "        outputs = model(inputs)\n",
    "\n",
    "        loss = criterion(outputs, targets)\n",
    "\n",
    "        # measure accuracy and record loss\n",
    "        prec1, prec2 = accuracy(outputs.data, targets.data, topk=(1, 1))        # prec1, prec5 = accuracy(outputs.data, targets.data, topk=(1, 5))\n",
    "        losses.update(loss.data, inputs.size(0))\n",
    "        top1.update(prec1, inputs.size(0))\n",
    "        top5.update(prec2, inputs.size(0))\n",
    "\n",
    "        # compute gradient and do SGD step\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        # measure elapsed time\n",
    "        batch_time.update(time.time() - end)\n",
    "        end = time.time()\n",
    "\n",
    "        # plot progress\n",
    "        bar.suffix  = '({batch}/{size}) Data: {data:.3f}s | Batch: {bt:.3f}s | Total: {total:} | ETA: {eta:} | Loss: {loss:.4f} | top1: {top1: .4f} | top5: {top5: .4f}'.format(\n",
    "                    batch=batch_idx + 1,\n",
    "                    size=len(train_loader),\n",
    "                    data=data_time.val,\n",
    "                    bt=batch_time.val,\n",
    "                    total=bar.elapsed_td,\n",
    "                    eta=bar.eta_td,\n",
    "                    loss=losses.avg,\n",
    "                    top1=top1.avg,\n",
    "                    top5=top5.avg,\n",
    "                    )\n",
    "        if (batch_idx) % show_step == 0:\n",
    "            print(bar.suffix)\n",
    "        bar.next()\n",
    "    bar.finish()\n",
    "    return (losses.avg, top1.avg)\n",
    "\n",
    "def test(val_loader, model, criterion, epoch, use_cuda):\n",
    "    global best_acc\n",
    "\n",
    "    batch_time = AverageMeter()\n",
    "    data_time = AverageMeter()\n",
    "    losses = AverageMeter()\n",
    "    top1 = AverageMeter()\n",
    "    top5 = AverageMeter()\n",
    "\n",
    "    # switch to evaluate mode\n",
    "    model.eval()\n",
    "    torch.set_grad_enabled(False)\n",
    "\n",
    "    end = time.time()\n",
    "    bar = Bar('Processing', max=len(val_loader))\n",
    "\n",
    "\n",
    "    for batch_idx, (inputs, targets) in enumerate(val_loader):\n",
    "        # measure data loading time\n",
    "        data_time.update(time.time() - end)\n",
    "\n",
    "        if use_cuda:\n",
    "            inputs, targets = inputs.cuda(), targets.cuda()\n",
    "        inputs, targets = torch.autograd.Variable(inputs, volatile=True), torch.autograd.Variable(targets)\n",
    "\n",
    "        # compute output\n",
    "        outputs = model(inputs)\n",
    "\n",
    "        loss = criterion(outputs, targets)\n",
    "\n",
    "        # measure accuracy and record loss\n",
    "        prec1, prec2 = accuracy(outputs.data, targets.data, topk=(1, 1))        #prec1, prec5 = accuracy(outputs.data, targets.data, topk=(1, 5))\n",
    "        # losses.update(loss.data[0], inputs.size(0))\n",
    "        losses.update(loss.data, inputs.size(0))\n",
    "        #top1.update(prec1[0], inputs.size(0))\n",
    "        top1.update(prec1, inputs.size(0))\n",
    "        #top5.update(prec5[0], inputs.size(0))\n",
    "        top5.update(prec2, inputs.size(0))\n",
    "\n",
    "        # measure elapsed time\n",
    "        batch_time.update(time.time() - end)\n",
    "        end = time.time()\n",
    "\n",
    "        # plot progress\n",
    "        bar.suffix  = '({batch}/{size}) Data: {data:.3f}s | Batch: {bt:.3f}s | Total: {total:} | ETA: {eta:} | Loss: {loss:.4f} | top1: {top1: .4f} | top5: {top5: .4f}'.format(\n",
    "                    batch=batch_idx + 1,\n",
    "                    size=len(val_loader),\n",
    "                    data=data_time.avg,\n",
    "                    bt=batch_time.avg,\n",
    "                    total=bar.elapsed_td,\n",
    "                    eta=bar.eta_td,\n",
    "                    loss=losses.avg,\n",
    "                    top1=top1.avg,\n",
    "                    top5=top5.avg,\n",
    "                    )\n",
    "        bar.next()\n",
    "    print(bar.suffix)\n",
    "    bar.finish()\n",
    "    return (losses.avg, top1.avg)\n",
    "\n",
    "def save_checkpoint(state, is_best, checkpoint='checkpoint', filename='checkpoint.pth.tar'):\n",
    "    filepath = os.path.join(checkpoint, filename)\n",
    "    torch.save(state, filepath)\n",
    "    if is_best:\n",
    "        shutil.copyfile(filepath, os.path.join(checkpoint, 'model_best.pth.tar'))\n",
    "\n",
    "def adjust_learning_rate(optimizer, epoch):\n",
    "    global state\n",
    "    if epoch in args.schedule:\n",
    "        state['lr'] *= args.gamma\n",
    "        for param_group in optimizer.param_groups:\n",
    "            param_group['lr'] = state['lr']\n",
    "\n",
    "def ToImage(tensor):\n",
    "    unloader = transforms.ToPILImage()\n",
    "    image = tensor.cpu().clone()\n",
    "    image = image.squeeze(0)\n",
    "    image = unloader(image)\n",
    "    image.show()\n",
    "    return image\n",
    "\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    main()\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "detectron2",
   "language": "python",
   "name": "detectron2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
